import mteb
from mteb.model_meta import ModelMeta
from colqwen2_wrappers.colqwen2_wrapper import ColQwen2Wrapper

BENCHMARKS = ["ViDoRe(v1)", "ViDoRe(v2)"]

# === 加载模型 ===
custom_model = custom_model_meta.load_model(
    MODEL_NAME,
    device="cuda",
    quantization=None
)

# === 加载任务并评估 ===
tasks = mteb.get_benchmarks(names=BENCHMARKS)
evaluator = mteb.MTEB(tasks=tasks)
results = evaluator.run(
    custom_model,
    output_folder="./test_output",
    batch_size=16
)

# === 打印结果 ===
for task_name, task_result in results.items():
    print(f"\nTask: {task_name}")
    print(f"Scores: {task_result['scores']}")